"""
多模态知识库管理器
整合OCR结果与文本知识，支持多格式内容索引和检索
"""

import os
import json
import asyncio
import hashlib
import base64
from typing import Dict, Any, List, Optional, Tuple, Union
from dataclasses import dataclass, field
from datetime import datetime
from PIL import Image
import io

from utils.logger import logger
from search_engine import SearchEngineManager, SearchDocument, search_engine_manager
from apps.agents.ocr_agent import OCRAgent


@dataclass
class MultimodalDocument:
    """多模态文档数据结构"""
    id: str
    title: str
    content: str
    content_type: str  # text, image, pdf, etc.
    source_url: str = ""
    source_path: str = ""
    metadata: Dict[str, Any] = field(default_factory=dict)
    ocr_result: Optional[Dict[str, Any]] = None
    embedding: Optional[List[float]] = None
    created_at: datetime = field(default_factory=datetime.now)
    updated_at: datetime = field(default_factory=datetime.now)
    tags: List[str] = field(default_factory=list)


class MultimodalKnowledgeBase:
    """多模态知识库管理器"""
    
    def __init__(self, kb_name: str = "default_kb"):
        self.kb_name = kb_name
        self.search_engine = search_engine_manager  # 复用现有的搜索引擎
        self.ocr_agent = OCRAgent()
        self.documents: Dict[str, MultimodalDocument] = {}
        self.index_dir = os.path.join("./data/knowledge_bases", kb_name)
        os.makedirs(self.index_dir, exist_ok=True)
        
        logger.info(f"多模态知识库初始化完成: {kb_name}")
    
    def _generate_doc_id(self, content: str, source: str = "") -> str:
        """生成文档ID"""
        combined = f"{content[:1000]}_{source}_{datetime.now().timestamp()}"
        return hashlib.md5(combined.encode()).hexdigest()
    
    async def add_text_document(self, title: str, content: str, source: str = "", 
                              metadata: Optional[Dict[str, Any]] = None, 
                              tags: Optional[List[str]] = None) -> str:
        """
        添加文本文档到知识库
        
        Args:
            title: 文档标题
            content: 文档内容
            source: 文档来源
            metadata: 文档元数据
            tags: 文档标签
            
        Returns:
            文档ID
        """
        doc_id = self._generate_doc_id(content, source)
        
        # 创建多模态文档
        multimodal_doc = MultimodalDocument(
            id=doc_id,
            title=title,
            content=content,
            content_type="text",
            source_url=source if source.startswith("http") else "",
            source_path=source if not source.startswith("http") else "",
            metadata=metadata or {},
            tags=tags or []
        )
        
        # 添加到搜索引擎索引
        search_doc = SearchDocument(
            id=doc_id,
            title=title,
            content=content,
            url=multimodal_doc.source_url or multimodal_doc.source_path or f"kb://{self.kb_name}/{doc_id}",
            doc_type="text",
            tags=tags or [],
            metadata={
                **multimodal_doc.metadata,
                "content_type": "text",
                "kb_name": self.kb_name
            }
        )
        
        if self.search_engine.add_document(search_doc):
            self.documents[doc_id] = multimodal_doc
            await self._persist_document(multimodal_doc)
            logger.info(f"文本文档添加到知识库: {title} ({doc_id})")
            return doc_id
        
        raise ValueError(f"添加文本文档失败: {title}")
    
    async def add_image_document(self, title: str, image_data: Union[str, bytes, Image.Image], 
                               source: str = "", perform_ocr: bool = True,
                               metadata: Optional[Dict[str, Any]] = None,
                               tags: Optional[List[str]] = None) -> str:
        """
        添加图像文档到知识库
        
        Args:
            title: 文档标题
            image_data: 图像数据 (base64字符串、字节或PIL Image)
            source: 图像来源
            perform_ocr: 是否执行OCR
            metadata: 文档元数据
            tags: 文档标签
            
        Returns:
            文档ID
        """
        # 处理图像数据
        if isinstance(image_data, str):
            # base64字符串
            image_bytes = base64.b64decode(image_data)
        elif isinstance(image_data, bytes):
            image_bytes = image_data
        else:
            # PIL Image
            img_byte_arr = io.BytesIO()
            image_data.save(img_byte_arr, format='PNG')
            image_bytes = img_byte_arr.getvalue()
        
        # 生成文档ID
        content_hash = hashlib.md5(image_bytes).hexdigest()
        doc_id = self._generate_doc_id(content_hash, source)
        
        # 执行OCR
        ocr_result = None
        ocr_text = ""
        if perform_ocr:
            try:
                task_data = {
                    "image_data": base64.b64encode(image_bytes).decode('utf-8'),
                    "detailed_info": True,
                    "clean_text": True
                }
                ocr_result = await self.ocr_agent.process_task(task_data)
                ocr_text = ocr_result.get("text", "")
                logger.info(f"图像OCR完成，识别到文本长度: {len(ocr_text)}")
            except Exception as e:
                logger.error(f"OCR处理失败: {str(e)}")
        
        # 创建多模态文档
        multimodal_doc = MultimodalDocument(
            id=doc_id,
            title=title,
            content=ocr_text,
            content_type="image",
            source_url=source if source.startswith("http") else "",
            source_path=source if not source.startswith("http") else "",
            metadata={
                **(metadata or {}),
                "image_size": len(image_bytes),
                "has_ocr": bool(perform_ocr)
            },
            ocr_result=ocr_result,
            tags=tags or []
        )
        
        # 添加到搜索引擎索引
        search_content = ocr_text if perform_ocr else f"[图像内容] {title}"
        search_doc = SearchDocument(
            id=doc_id,
            title=title,
            content=search_content,
            url=multimodal_doc.source_url or multimodal_doc.source_path or f"kb://{self.kb_name}/{doc_id}",
            doc_type="image",
            tags=["image"] + (tags or []),
            metadata={
                **multimodal_doc.metadata,
                "content_type": "image",
                "kb_name": self.kb_name,
                "has_ocr": bool(perform_ocr)
            }
        )
        
        if self.search_engine.add_document(search_doc):
            self.documents[doc_id] = multimodal_doc
            await self._persist_document(multimodal_doc, image_bytes if perform_ocr else None)
            logger.info(f"图像文档添加到知识库: {title} ({doc_id})")
            return doc_id
        
        raise ValueError(f"添加图像文档失败: {title}")
    
    async def _persist_document(self, document: MultimodalDocument, image_bytes: Optional[bytes] = None):
        """持久化文档"""
        doc_dir = os.path.join(self.index_dir, document.id[:2], document.id[2:4])
        os.makedirs(doc_dir, exist_ok=True)
        
        # 保存文档元数据
        doc_path = os.path.join(doc_dir, f"{document.id}.json")
        with open(doc_path, 'w', encoding='utf-8') as f:
            doc_dict = {
                **document.__dict__,
                "created_at": document.created_at.isoformat(),
                "updated_at": document.updated_at.isoformat()
            }
            json.dump(doc_dict, f, ensure_ascii=False, indent=2)
        
        # 如果有图像数据，保存图像
        if image_bytes and document.content_type == "image":
            image_path = os.path.join(doc_dir, f"{document.id}.png")
            with open(image_path, 'wb') as f:
                f.write(image_bytes)
    
    def search(self, query: str, limit: int = 10, 
              content_type: Optional[str] = None, 
              tags: Optional[List[str]] = None) -> List[MultimodalDocument]:
        """
        搜索知识库
        
        Args:
            query: 搜索查询
            limit: 返回结果数量限制
            content_type: 内容类型过滤
            tags: 标签过滤
            
        Returns:
            匹配的多模态文档列表
        """
        search_result = self.search_engine.search_documents(
            query=query,
            limit=limit,
            doc_type=content_type,
            tags=tags
        )
        
        # 转换为多模态文档
        results = []
        for search_doc in search_result.documents:
            # 先从内存缓存中查找
            if search_doc.id in self.documents:
                results.append(self.documents[search_doc.id])
            else:
                # 尝试从持久化存储中加载
                doc = self._load_document_from_disk(search_doc.id)
                if doc:
                    results.append(doc)
        
        return results
    
    def _load_document_from_disk(self, doc_id: str) -> Optional[MultimodalDocument]:
        """从磁盘加载文档"""
        doc_dir = os.path.join(self.index_dir, doc_id[:2], doc_id[2:4])
        doc_path = os.path.join(doc_dir, f"{doc_id}.json")
        
        try:
            if os.path.exists(doc_path):
                with open(doc_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                
                # 恢复日期对象
                data["created_at"] = datetime.fromisoformat(data["created_at"])
                data["updated_at"] = datetime.fromisoformat(data["updated_at"])
                
                doc = MultimodalDocument(**data)
                self.documents[doc_id] = doc
                return doc
        except Exception as e:
            logger.error(f"加载文档失败: {doc_id}, 错误: {str(e)}")
        
        return None
    
    async def batch_add_documents(self, documents: List[Dict[str, Any]]) -> List[Tuple[str, bool]]:
        """
        批量添加文档
        
        Args:
            documents: 文档列表
            
        Returns:
            [(文档ID, 是否成功)] 列表
        """
        results = []
        
        for doc_data in documents:
            try:
                if doc_data.get("content_type") == "image":
                    doc_id = await self.add_image_document(
                        title=doc_data["title"],
                        image_data=doc_data["content"],
                        source=doc_data.get("source", ""),
                        perform_ocr=doc_data.get("perform_ocr", True),
                        metadata=doc_data.get("metadata"),
                        tags=doc_data.get("tags")
                    )
                else:
                    doc_id = await self.add_text_document(
                        title=doc_data["title"],
                        content=doc_data["content"],
                        source=doc_data.get("source", ""),
                        metadata=doc_data.get("metadata"),
                        tags=doc_data.get("tags")
                    )
                results.append((doc_id, True))
            except Exception as e:
                logger.error(f"批量添加文档失败: {str(e)}")
                results.append(("", False))
        
        return results
    
    def delete_document(self, doc_id: str) -> bool:
        """
        删除文档
        
        Args:
            doc_id: 文档ID
            
        Returns:
            是否成功删除
        """
        if self.search_engine.delete_document(doc_id):
            if doc_id in self.documents:
                del self.documents[doc_id]
            
            # 删除持久化文件
            try:
                doc_dir = os.path.join(self.index_dir, doc_id[:2], doc_id[2:4])
                doc_path = os.path.join(doc_dir, f"{doc_id}.json")
                if os.path.exists(doc_path):
                    os.remove(doc_path)
                
                image_path = os.path.join(doc_dir, f"{doc_id}.png")
                if os.path.exists(image_path):
                    os.remove(image_path)
            except Exception as e:
                logger.error(f"删除文档文件失败: {str(e)}")
            
            logger.info(f"文档从知识库删除: {doc_id}")
            return True
        
        return False
    
    def get_document(self, doc_id: str) -> Optional[MultimodalDocument]:
        """
        获取文档
        
        Args:
            doc_id: 文档ID
            
        Returns:
            多模态文档或None
        """
        # 先从内存缓存查找
        if doc_id in self.documents:
            return self.documents[doc_id]
        
        # 从磁盘加载
        return self._load_document_from_disk(doc_id)
    
    def get_kb_stats(self) -> Dict[str, Any]:
        """
        获取知识库统计信息
        
        Returns:
            统计信息
        """
        search_stats = self.search_engine.get_search_stats()
        
        # 统计内容类型分布
        type_count = {"text": 0, "image": 0, "other": 0}
        for doc in self.documents.values():
            if doc.content_type in type_count:
                type_count[doc.content_type] += 1
            else:
                type_count["other"] += 1
        
        return {
            "kb_name": self.kb_name,
            "total_documents": len(self.documents),
            "type_distribution": type_count,
            "search_index_stats": search_stats,
            "index_dir": self.index_dir
        }


# 创建全局多模态知识库实例
multimodal_kb = MultimodalKnowledgeBase()


if __name__ == "__main__":
    # 测试多模态知识库
    import asyncio
    
    async def test_multimodal_kb():
        """测试多模态知识库"""
        kb = MultimodalKnowledgeBase("test_kb")
        
        # 添加文本文档
        text_doc_id = await kb.add_text_document(
            title="Python基础知识",
            content="Python是一种解释型、面向对象、动态数据类型的高级程序设计语言。Python语法简洁清晰，特色之一是强制用空白符作为语句缩进。",
            source="local",
            tags=["programming", "python", "basics"]
        )
        print(f"添加文本文档: {text_doc_id}")
        
        # 模拟添加图像文档（需要实际图像数据才能测试完整功能）
        print("提示：要测试图像功能，请提供实际图像数据")
        
        # 搜索测试
        print("\n搜索测试:")
        results = kb.search("Python 语法", limit=5)
        for i, doc in enumerate(results, 1):
            print(f"{i}. {doc.title} (类型: {doc.content_type})")
            print(f"   内容: {doc.content[:100]}...")
        
        # 统计信息
        stats = kb.get_kb_stats()
        print("\n知识库统计:")
        print(json.dumps(stats, indent=2, ensure_ascii=False, default=str))
    
    # 运行测试
    asyncio.run(test_multimodal_kb())
